from kivy.metrics import dp, sp
from kivymd.uix.boxlayout import MDBoxLayout
from kivymd.uix.button import MDRoundFlatButton, MDIconButton
from kivymd.uix.filemanager import MDFileManager
from kivymd.uix.label import MDLabel
from kivymd.uix.gridlayout import GridLayout
from kivymd.uix.list import MDList
from kivymd.uix.progressbar import MDProgressBar
from kivymd.uix.snackbar import MDSnackbar
from kivy.core.window import Window
from kivymd.app import MDApp
from ultralytics import YOLO

KV = '''
MDBoxLayout:
    MDProgressBar:
        id: progress
        pos_hint: {"center_x": .5, "center_y": .6}
        type: "indeterminate"
'''

class MyApp(MDApp):
    def build(self):
        Window.bind(on_dropfile=self.on_dropfile)
        main_layout = MDBoxLayout(orientation='vertical', spacing=10, padding=10)

        # Создаем метку для инструкции
        instruction_label = MDLabel(text='Перетащите файл на изображение или выберите с помощью кнопки', height=50,
                                    halign='center')

        # Создаем кнопку для загрузки файла
        button = MDRoundFlatButton(text='Загрузить файл(ы)')
        button.bind(on_press=self.load_file)

        # Добавляем виджеты на экран
        main_layout.add_widget(instruction_label)

        main_layout.add_widget(button)

        return main_layout


    def load_file(self, instance):
        self.file_manager = MDFileManager(preview=True, icon_folder='image/folder.png', exit_manager=self.exit_manager,
                                          select_path=self.select_path)
        self.file_manager.show("C:/")
        model = YOLO('best2.pt')
        p1 = file_path.decode('utf-8')
        inference_result = model.predict(p1, conf=0.2, save=False)
        for i in inference_result:
            print(i.masks)

    def select_path(self, path: str):
        self.exit_manager()
        MDSnackbar(
            MDLabel(
                text=path,
            ),
            y=dp(24),
            pos_hint={"center_x": 0.5},
            size_hint_x=0.8,
        ).open()
        result_layout = GridLayout(cols=3)
        for i in range(3):
            box = MDList(spacing="8dp")
            folder_button = MDIconButton(icon='image/folder_w.png', icon_size=sp(100),)
            folder_name = ["Олень", "Косуля", "Кабарга"]
            folder_button.bind(on_press=lambda _, folder=folder_name: self.open_folder(folder))
            text_folder = MDLabel(text=folder_name[i], halign='left', padding=sp(32))
            box.add_widget(folder_button)
            box.add_widget(text_folder)
            result_layout.add_widget(box)

        self.root.clear_widgets()
        self.root.add_widget(result_layout)

    def exit_manager(self, *args):
        self.manager_open = False
        self.file_manager.close()

    def on_dropfile(self, window, file_path):
        main_pb = MDBoxLayout(orientation='vertical', spacing=5, padding=400)
        pb = MDProgressBar(pos_hint={"center_x": 0.5, "center_y": .1}, type="indeterminate", size_hint_max_x=700)
        pb.check_size(self)
        pb.start()
        main_pb.add_widget(pb)
        self.root.clear_widgets()
        self.root.add_widget(main_pb)
        model = YOLO('best2.pt')
        p1 = file_path.decode('utf-8')
        inference_result = model.predict(p1, conf=0.2, save=False)
        ans = [1]*len(inference_result)

        for i in inference_result:
        masks = i.xy
        labels = i.labels
        if i.labels == 'SR':
            return 1
        elif i.labels == 'BR':
            return 0
        elif i.labels == 'Z':
            return 2
        else:
        return 1
        self.root.clear_widgets()

        result_layout = GridLayout(cols=3)
        for i in range(3):
            box = MDList(spacing="8dp")
            folder_button = MDIconButton(icon='image/folder_w.png', icon_size=sp(100), )
            folder_name = ["Олень", "Косуля", "Кабарга"]
            folder_button.bind(on_press=lambda _, folder=folder_name: self.open_folder(folder))
            text_folder = MDLabel(text=folder_name[i], halign='left', padding=sp(32))
            box.add_widget(folder_button)
            box.add_widget(text_folder)
            result_layout.add_widget(box)

        self.root.clear_widgets()
        self.root.add_widget(result_layout)


    def open_folder(self, folder_name):
        pass


MyApp().run()
